# 取股票ID

import cv2
import numpy as np
import pytesseract
import time
from shape_judgment.screen_region_storage import get_region

# 参数: cv2.imread 这个函数的返回值可以保存
def save_pixel_array_to_png(pixel_array, output_path):
    """
    将像素数组（BGRA格式）保存为PNG图片
    :param pixel_array: 由screen_region_to_pixel_array返回的像素数组（shape: (h, w, 4)）
    :param output_path: 输出PNG文件路径（如"sub_region.png"）
    """
    if pixel_array is None:
        raise ValueError("像素数组为空，无法保存")
    
    # 校验数组格式（必须是BGRA四通道）
    if len(pixel_array.shape) != 3 or pixel_array.shape[2] != 4:
        raise ValueError(f"像素数组格式错误，需为BGRA四通道，实际为{pixel_array.shape}")
    
    # 直接使用OpenCV保存（PNG格式会保留alpha通道）
    # cv2.imwrite默认支持BGRA格式保存为PNG
    success = cv2.imwrite(output_path, pixel_array)
    
    if not success:
        raise IOError(f"保存PNG失败，请检查路径是否有效：{output_path}")
    print(f"PNG图片已保存至：{output_path}")


# 用于截取图像, 用于调试
def screen_region_to_pixel_array(x, y, width, height, screen_img):
    """
    从已有的股票软件截图中截取指定子区域，并转换为像素点数组（含alpha通道）
    :param x: 子区域左上角X坐标（相对于screen_img的局部坐标）
    :param y: 子区域左上角Y坐标（相对于screen_img的局部坐标）
    :param width: 子区域宽度
    :param height: 子区域高度
    :param screen_img: 已截取的股票软件全屏截图（BGR格式，OpenCV图像）
    :return: 子区域像素数组，shape为(height, width, 4)，每个元素是(B, G, R, A)
    """
    # 1. 校验输入截图有效性
    if screen_img is None or screen_img.size == 0:
        raise ValueError("输入的股票软件截图为空")
    # 2. 校验坐标是否在截图范围内（防止越界）    
    screen_height, screen_width = screen_img.shape[:2]
    if (x < 0 or y < 0 or 
        x + width > screen_width or 
        y + height > screen_height):
        # save_frame_to_image(screen_img,"ERR 144.png")
        print(f"!!!!    51 截取区域超出截图范围！截图尺寸: {screen_width}×{screen_height}, 截取区域: ({x}, {y})-{x+width}×{y+height}")
        # raise ValueError(
        #     f"截取区域超出截图范围！截图尺寸: {screen_width}×{screen_height}, "
        #     f"截取区域: ({x}, {y})-{x+width}×{y+height}"
        # )
        return None
    
    # 3. 从全屏截图中截取子区域（BGR格式）
    sub_region = screen_img[y:y+height, x:x+width].copy()
    
    # 4. 转换为BGRA格式（添加alpha通道，默认完全不透明）
    if sub_region.shape[2] == 3:  # 若输入是BGR格式（无alpha）
        bgra_region = cv2.cvtColor(sub_region, cv2.COLOR_BGR2BGRA)
        bgra_region[:, :, 3] = 255  # alpha通道设为255（完全不透明）
    else:  # 若输入已含alpha通道（罕见）
        bgra_region = sub_region
    
    return bgra_region



# 可能会有多个匹配, 返回所有匹配先
def find_all_matches(screen_img, template_img, threshold):
    """获取所有符合阈值的匹配结果，返回列表[(x, y, w, h, 匹配度)]"""
    # # 预处理（使用之前优化的预处理函数）
    # screen_processed = preprocess_image(screen_img)
    # template_processed = preprocess_image(template_img, is_template=True)
    h, w = template_img.shape[:2]
    
    if h == 0 or w == 0:
        return []
    
    if screen_img.size < template_img.size:  # 截图尺寸小于模板尺寸
        print("截图尺寸小于模板尺寸")
        return []
    
    # 执行模板匹配并获取匹配度矩阵
    result = cv2.matchTemplate(
        screen_img, 
        template_img, 
        cv2.TM_CCOEFF_NORMED
    )
    # 找到所有超过阈值的位置
    locations = np.where(result >= threshold)
    matches = []
    
    # 遍历所有匹配位置，记录坐标和匹配度
    for y, x in zip(locations[0], locations[1]):
        match_value = result[y, x]  # 该位置的匹配度（0-1）
        matches.append((x, y, w, h, match_value))  # 存储x坐标、y坐标、宽、高、匹配度
    return matches

# 找 模板匹配
def find_template_position(screen_img, template_img, threshold=0.85):
    """
    在屏幕截图中查找模板位置
    :param screen_img: 全屏截图（BGR格式）
    :param template_img: 模板图片（BGR格式）
    :param threshold: 匹配阈值（0-1，越高越严格） 数值为0.8时,还是会有错误, 比如需匹配0.00%,匹配到的却是0.05%
    :return: 匹配区域的左上角坐标(x, y)，若未找到返回None
    """
    # # 获取模板尺寸
    # h, w = template_img.shape[:2]
    
    # # 执行模板匹配
    # result = cv2.matchTemplate(screen_img, template_img, cv2.TM_CCOEFF_NORMED)
    
    # # 查找匹配度超过阈值的位置
    # locations = np.where(result >= threshold)

    # # 取匹配度最高的位置
    # if len(locations[0]) > 0:
    #     # 找到最大匹配值的索引
    #     print("#多个匹配的情况")
    #     max_idx = np.argmax(result)
    #     h_idx, w_idx = np.unravel_index(max_idx, result.shape)
    #     return (w_idx, h_idx, w, h)  # (x, y, 宽, 高)

    all_matches  = find_all_matches(screen_img, template_img,threshold)
    if not all_matches:
        return None  # 无匹配结果
    # 按X坐标升序排序（X越小越靠左），取第一个
    # 若X相同，可选匹配度最高的
    all_matches.sort(key=lambda m: (-m[4]))  # 匹配度倒序排序
    rightmost = all_matches[0]
    return (rightmost[0], rightmost[1], rightmost[2], rightmost[3])  # 返回(x, y, w, h)
    # # 取第一个匹配位置（最相似的）
    # for pt in zip(*locations[::-1]):  # pt为左上角坐标
    #     print("#只有一个匹配的情况")
    #     return (pt[0], pt[1], w, h)  # 返回(x, y, 宽度, 高度)
    
    # return None  # 未找到匹配区域

def extract_numbers_from_image(img):
    # 使用Tesseract识别文本，添加参数提高识别准确性
    text = ''
    try:
        # --psm 7 表示假设图像是单行文本
        text = pytesseract.image_to_string(img, lang='eng', config='--psm 7 -c tessedit_char_whitelist=0123456789')
    except Exception as e:
        print(f"!!!!    ERR 151 提取文本时发生错误：{str(e)}")
    
    # 过滤出所有数字
    numbers = ''.join([c for c in text if c.isdigit()])
    
    return numbers

# 找出数字与中文名称之间的分隔
def compare_pixel_arrays(img_clip_array_2):
    # 这里黑色是(7,7,7)

    # 0 第一组全黑的开始
    # 1 第一组全黑结束
    # 2 找到第二组全黑的开始
    state = 0
    # 连续全黑列数 连续5组, 就是数字与中文名称的分隔区间
    number_of_all_black_columns = 0
    # 数字开始的坐标, 用于节省tesseract的工作量
    number_of_first_pos = 0

    # BEGIN 遍历图像中的像素点:
    height, width = img_clip_array_2.shape[:2]
    for x in range(width):
        matched_pixels_this_column = 0 
        for y in range(height):
            # 获取模板像素的BGRA值
            t_b, t_g, t_r, t_a = img_clip_array_2[y, x] 
            if 7==t_b and 7==t_g and 7==t_r:
                matched_pixels_this_column += 1
        if matched_pixels_this_column ==height:
            if state == 1:
                number_of_all_black_columns +=1
            if number_of_all_black_columns >=5:
                # return (number_of_first_pos,x-5)
                state = 2
        elif matched_pixels_this_column != height:
            if state == 0:
                state = 1
                number_of_first_pos = x
            if state == 2:
                return (number_of_first_pos,x-5)
    return (0,0)


def pixel_point_traversal_generate_numbers(img):

    # 遍历图像中每列:
    height, width = img.shape[:2]
    consecutive_empty_column = 0 # 连续空列数 连空三列就是结束
    result_str = "" # 结果字符串 是一个由数字组成的股票ID
    x = 0
    while x < width:
        # print("# 遍历列: " + str(x))
        matched_pixels_this_column = 0
        state_empty = True # 空的状态, 可能是空行
        for y in range(height):
            # 获取模板像素的BGRA值
            t_b, t_g, t_r, t_a = img[y, x] 
            if state_empty:
                if 7==t_b and 7==t_g and 7==t_r:
                    matched_pixels_this_column += 1
                    if(matched_pixels_this_column>=height):
                        # 是个空行
                        consecutive_empty_column +=1
                        if consecutive_empty_column >=3 and len(result_str):
                            # print("# 218 到达三条空列 退出, 股票ID: " + result_str +"\theight:",height,"\tmatched_pixels_this_column:",matched_pixels_this_column)
                            return result_str
                else:
                    state_empty = False
                    consecutive_empty_column = 0
            # 为应对7这个数字, 不能等第二轮循环, 否则错过第零个像素.
            if False == state_empty:
                if 0== y:
                    # 判断为7
                    if (t_r,t_g,t_b) == (6,83,144):
                        result_str += "7"
                        x+=8 # 字宽是9 留个1给for去加
                        break
                elif 1 == y:
                    # 判断为2
                    if (t_r,t_g,t_b) == (6,6,48):
                        result_str += "2"
                        x+=8 # 字宽是9 留个1给for去加
                        break
                elif 2 == y:
                    # 判断为 0、3、6、8、9
                    # 判断为 0、6、8:
                    if (t_r,t_g,t_b) == (6,6,48):
                        t_b_2,t_g_2,t_r_2 ,t_a = img[y+1, x]
                        if (6,48,114) == (t_r_2,t_g_2,t_b_2):
                            # 至此, 可判断为0、6
                            t_b_3,t_g_3,t_r_3 ,t_a= img[y+2, x]
                            if(6,48,114)==(t_r_3,t_g_3,t_b_3):
                                result_str += "0"
                            else:
                                result_str += "6"
                        elif (6,6,48)==(t_r_2,t_g_2,t_b_2):
                            result_str += "8"
                    # 判断为 3、9
                    elif (6,6,83) == (t_r,t_g,t_b) :
                        t_b_2,t_g_2,t_r_2,t_a = img[y+1, x]
                        if 7 == t_b_2:
                            result_str += "3"
                        else:
                            result_str += "9"
                    x+=8
                    break
                elif 3 == y:
                    # 判断为 1:
                    if (6,83,144) ==(t_r,t_g,t_b) :
                        result_str += "1"
                        x+=7
                        break
                elif 5 == y:
                    # 判断为5:
                    if (6,6,83)==(t_r,t_g,t_b) :
                        result_str += "5"
                        x+=8
                        break
                elif 7 == y:
                    # 判断为4:
                    if (6,48,114) == (t_r,t_g,t_b) :
                        result_str += "4"
                        x+=8
                        break
        x+=1


        # if matched_pixels_this_column ==height:
        #     return x-5

    return result_str


def pixel_point_traversal_generate_numbers_777(img):
    # 遍历图像中每列:
    height, width = img.shape[:2]
    consecutive_empty_column = 0 # 连续空列数 连空三列就是结束
    result_str = "" # 结果字符串 是一个由数字组成的股票ID
    x = 0
    while x < width:
        # print(f"# 293 x:{x}\tmatched_pixels_this_column:{consecutive_empty_column}")
        matched_pixels_this_column = 0
        state_empty = True # 空的状态, 可能是空行
        for y in [0,1,2,3,5,7]:
            # print(f" # 296 y:{y}")
            # 获取模板像素的BGRA值
            t_b, t_g, t_r, t_a = img[y, x] 
            if state_empty:
                if 7==t_b and 7==t_g and 7==t_r:
                    matched_pixels_this_column += 1
                    # print(f"# 303 matched_pixels_this_column += 1    == {matched_pixels_this_column} \t height:{height}")
                    if(matched_pixels_this_column>=6):
                        # 是个空行
                        consecutive_empty_column +=1
                        if consecutive_empty_column >=3 and len(result_str):
                            # print("# 218 到达三条空列 退出, 股票ID: " + result_str +"\theight:",height,"\tmatched_pixels_this_column:",matched_pixels_this_column)
                            return result_str
                else:
                    state_empty = False
                    consecutive_empty_column = 0
            # 为应对7这个数字, 不能等第二轮循环, 否则错过第零个像素.
            
            if 0 == y: # 判断为7                    
                if (t_r,t_g,t_b) == (7,71,136):
                    result_str += "7"
                    x+=8 # 字宽是9 留个1给for去加
                    break
            elif 1 == y: # 判断为2
                if (t_r,t_g,t_b) == (7,7,39):
                    result_str += "2"
                    x+=8
                    break
            elif 2 == y: # 判断为 0、3、6、8、9
                # print("# 324 0、3、6、8、9")
                # 0、6、8
                if (t_r,t_g,t_b) == (7,7,39):
                    # print("# 327 0、6、8")
                    t_b_2,t_g_2,t_r_2 ,t_a = img[y+1, x]
                    if (7,7,39) == (t_r_2,t_g_2,t_b_2):
                        result_str += "8"
                    else:
                        t_b_3,t_g_3,t_r_3 ,t_a = img[0, x+3]
                        if (168,200,200) == (t_r_3,t_g_3,t_b_3):
                            result_str+="6"
                        else:
                            # print(f"# 336 {(t_r_3,t_g_3,t_b_3)}  y:{y}  x:{x}")
                            result_str+="0"
                    x +=8
                    break
                # 3、9
                elif (t_r,t_g,t_b) == (7,7,71):
                    # print("# 341 3、9")
                    t_b_2,t_g_2,t_r_2,t_a = img[y+1, x]
                    if (7,71,136) == (t_r_2,t_g_2,t_b_2):
                        result_str += "9"
                    else:
                        result_str += "3"
                    x +=8
                    break
            elif 3 == y: # 1
                if (7,71,136) == (t_r,t_g,t_b):
                    result_str += "1"
                    x+=7
                    break
            elif 5 == y: # 5
                if (7,7,71) == (t_r,t_g,t_b):
                    result_str += "5"
                    x+=8
                    break
            elif 7 == y: # 4
                if (7,39,104) == (t_r,t_g,t_b):
                    result_str += "4"
                    x+=8
                    break
            # END else
        # END for
        x+=1
    # END while
    return result_str



global_debug_number = 0

# 取股票ID
def get_stock_id(img):
    # 第一个参数: OpenCV的截图类别
    # 返回值: 股票ID 是这字符串
    # 之前的屏幕截图函数是: cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR) 用cv2这类函数的返回值才能正常保存成图片
    # 程序中, 通过进一步裁剪多余的区域, 以达到更小的计算量 一行标题栏, 一行菜单栏, 共56像素. 所以Y的起始定为56
    # 为什么要判断两个三角的位置, 因为这个区域可能是被折叠的, 看不见, 这个可折叠侧边栏共有三种状态, 所以, 判断有没有三角按钮是取到股票ID的关键.
    # 取到三角按钮后, 再进一步截取更小的范围, 作进一步处理.

    time1 = time.time()

    # 2. 在剪裁区域中查找三角按钮
    # 2.1 加载模板图片
    template_img_left = cv2.imread("eastmoney desktop software picture/template_triangle_button_left.png")
    template_img_right = cv2.imread("eastmoney desktop software picture/template_triangle_button_right.png")
    # print ("调用OpenCV模板比对")
    match_pos1 = find_template_position(img, template_img_left)
    match_pos2 = find_template_position(img, template_img_right)
    # pos1与pos2, 都包含坐标与大小(np.int64(69), np.int64(10), 8, 15)(np.int64(265), np.int64(10), 8, 15)

    if match_pos1 is None or match_pos2 is None:
        return None

    # 3 再次截取
    # 3.1 取起止坐标
    clip_2_pos_1 = (match_pos1[0]+match_pos1[2],match_pos1[1]) #+match_pos1[3]
    clip_2_pos_2 = (match_pos2[0],match_pos2[1]+match_pos2[3])
    # # 3.2 截取:
    # clip_2 = cropped[clip_2_pos_1[1]:clip_2_pos_2[1], clip_2_pos_1[0]:clip_2_pos_2[0]]
    # 3.3 调试:
    # print(f"#176     {clip_2_pos_1}\t{clip_2_pos_2}")
    img_clip_array_2 = screen_region_to_pixel_array(clip_2_pos_1[0],clip_2_pos_1[1]
                                                ,clip_2_pos_2[0]-clip_2_pos_1[0],clip_2_pos_2[1]-clip_2_pos_1[1],img)
    if img_clip_array_2 is None :
        return None

    # 数字高为12像素
    img_clip_array_4 = img_clip_array_2[1:1+12, 1:-1].copy()
    stock_numbers = pixel_point_traversal_generate_numbers(img_clip_array_4)
    if 6 > len(stock_numbers): # 6 < 7
        stock_numbers = pixel_point_traversal_generate_numbers_777(img_clip_array_4)
    # print("# 369 获取股票ID:",stock_numbers)
    # time2 = time.time()
    # print("# 372 优化前用时:",time2-time1) # 优化前用时:0.030946016311645508 # 优化后用时: 0.0017037391662597656

    if 0 == len(stock_numbers):
        template_img_689009 = cv2.imread("eastmoney desktop software picture/689009.png")
        match_pos_336 = find_template_position(img, template_img_689009, .95)
        if match_pos_336 is None:
            global global_debug_number
            # cv2.imwrite(f"get_stock_id_ERR_335_{str(global_debug_number)}_0.png", img)
            # cv2.imwrite(f"get_stock_id_ERR_335_{str(global_debug_number)}_1.png", img_clip_array_2)
            # cv2.imwrite(f"get_stock_id_ERR_335_{str(global_debug_number)}_2.png", img_clip_array_4)
            global_debug_number+=1
            return '`'
        else:
            return '`689009'

    return '`'+stock_numbers

def get_stock_id_334(screen_shot_img):
    # screen_shot_img 整个窗口
    # return stock_id 是一个字符串, 如果没取到, 就返回None

    window_height,window_width = screen_shot_img.shape[:2]
    # 取,原本侧边栏1的股票名称区域坐标:
    _x1,_x2,_y1,_y2 = get_region("股票名称与价格")
    # 1. 计算小区域
    delta_x = window_width-_x2-2
    _x1 = _x1+delta_x
    _x2 = _x2+delta_x
    # 2. 从小区域中, 截取图像
    img_stock_name = screen_shot_img[_y1:_y2,_x1:_x2]
    # 3. 从小区域中, 取得股票ID

    return get_stock_id(img_stock_name)


    return None


if __name__ == "__main__":
    # 调试:
    print("452")


# 数字像素点颜色特征
#   最多判断到第一列第4个像素就可以判断出是哪个数字,除了1之后, 其它数字都是9列
#   0: 2 (6,6,48)	与6,8冲突
#   	3(6,48,114)
#   	4(6,48,114)
#   1: 3 (6,83,144)
#   2: 1 (6,6,48)	
#   3: 2 (6,6,83)	与9冲突
#   	3(7,7,7)
#   4: 7 (6,48,114)
#   5: 5 (6,6,83)
#   6: 2 (6,6,48)	有冲突
#   	3(6,48,114) 还是与零有冲突
#   	4(6,83,144) 
#   7: 0 (6,83,144)
#   8: 2 (6,6,48)	有冲突
#   	3(6,6,48)
#   9: 2 (6,6,83)	与3冲突
#   	3(6,83,144)
#   
# 规则:1这个数字, 在它前面必须空一行, 后面空两行, 但如果后面还是1,那么后面这个1的前面一行空行, 与之前一个1的后面空行是重叠的.
#   数字之间的空行最多是两行, 所以,判断到第三行空行, 就结束判断.
# 

# 现在发现 689009这只股票只有窗口缩小后会变两行, 而侧边栏在窗口缩小后, 宽度会变小
#    而且, 分成两行后, 字体变小了.
#    超4个字的股票名极少, 所以暂用图像比对就行.

# BUG 不知为何, 会有取不到股票ID的情况, 取到了只有`
#       增加了一些调试语句 输出 f"get_stock_id_ERR_335_{str(global_debug_number)}.png"
# 原本(6,6,48)这样的颜色现在变成(7,7,39) 图像的裁剪没问题, 问题出在文字的解析. 原本的黑底是(6,6,6),现在变成(7,7,7)
# 用tesseract解析不可行, 可能会有认错数字的情况. 而且速度还慢.
# 另一套数字分析: 现在了解了不用Y轴上每个像素遍历, 遍历关键的几个就行
# 担心会混掉, 所以独立一个函数出来, 如果用到 or 怕是会解析出错的数字.
# 
# 7: 0(7,71,136)
# 2: 1(7,7,39)
# 6: 2(7,7,39)
# 0: 2(7,7,39)
# 8: 2(7,7,39)
# 3: 2(7,7,71)
# 9: 2(7,7,71)
# 1: 3(7,71,136)
# 5: 5(7,7,71)
# 4: 7(7,39,104)
# 用于判断的像素点[0,1,2,3,5,7]

